from keras.datasets import cifar10
from keras.preprocessing.image import ImageDataGenerator
from keras.models import Sequential
from keras.layers.core import Dense, Dropout, Activation, Flatten
from keras.layers.convolutional import Convolution2D, MaxPooling2D
from keras.optimizers import SGD,adam
from keras.utils import np_utils
from Data_utils.Load import load_train17,load_val17,load_train32,load_val32,load_train18,load_val18,load_train32C,load_val32C
import numpy as np
from Data_utils.Crrupt import darken,add_gaussian_noise, crop_center_board_6, crop_center_board_3, crop_center_board
from model import baseline_conv, up_pooling_model, linear_conv, SRCNN, SRCNN_test, baseline_conv_test,linear_conv_test,residual_conv,residual_conv_test
import os
from skimage import io,util,exposure,measure
import skimage
import matplotlib.pyplot as plt
import sys
import keras.backend as K
from Data_utils.Metrics import PSNR
import json
import tensorflow as tf
import numpy

def custom_loss(y_true,y_pred):
    MSE = K.mean(K.square(y_pred - y_true))
    constraint = K.mean(K.abs(y_pred))
    loss = 5*MSE + constraint
    return loss

def PSNR_loss(I,P):
    """
    #sess = K.get_session()
    print 'I: {0}'.format(I)
    print 'I shape: {0}'.format(tf.shape(I))
    print 'I size: {0}'.format(tf.size(I))
    """
    batch_size = 128.
    I = K.reshape(I,(int(batch_size),18*18))
    P = K.reshape(P,(int(batch_size),18*18))
    MSE = K.mean(K.square(I-P),axis=1)
    PSNR_ = 20. * (tf.log(255./K.sqrt(MSE))/tf.log(10.))
    PSNR_ = K.mean(PSNR_)
    return -PSNR_

def train_model():
    train_imgs, train_names = load_train32C(Debug=False)
    val_imgs, val_names = load_val32C(Debug=False)
    train_Y = np.asarray(train_imgs)
    train_X = darken(train_Y)
    #train_Y = crop_center_board(train_Y,3)
    #train_X = add_gaussian_noise(train_X)
    val_Y = np.asarray(val_imgs)
    val_X = darken(val_Y)
    #val_Y = crop_center_board(val_Y,3)
    #val_X = add_gaussian_noise(val_X)
    print '{0} train samples'.format(len(train_Y))
    print '{0} val samples'.format(len(val_Y))
    print 'patch size: {0}'.format(train_Y[0].shape[0])
    """
    train_X = np.reshape(train_X,(train_X.shape[0],1,train_X[0].shape[0],train_X[0].shape[1]))
    train_Y = np.reshape(train_Y,(train_Y.shape[0],1,train_Y[0].shape[0],train_Y[0].shape[1]))
    train_X /= 255.
    train_Y /= 255.
    val_X = np.reshape(val_X,(val_X.shape[0],1,val_X[0].shape[0],val_X[0].shape[1]))
    val_Y = np.reshape(val_Y,(val_Y.shape[0],1,val_Y[0].shape[0],val_Y[0].shape[1]))
    val_X /= 255.
    val_Y /= 255.
    """
    train_X = np.transpose(train_X,(0,3,1,2))
    train_Y = np.transpose(train_Y,(0,3,1,2))
    train_X /= 255.
    train_Y /= 255.
    val_X = np.transpose(val_X,(0,3,1,2))
    val_Y = np.transpose(val_Y,(0,3,1,2))
    val_X /= 255.
    val_Y /= 255.
    with K.tf.device('/gpu:1'):
        K.set_session(K.tf.Session(config=K.tf.ConfigProto(allow_soft_placement=True, log_device_placement=True)))
        model = SRCNN(img_channels=3,img_rows=32,img_cols=32)
        model.compile(optimizer=SGD(lr=0.01,momentum=0.9,nesterov=True),loss= 'mse')
    """
    hist = model.fit({'conv_input': train_X, 'skip_input': train_X, 'output': train_Y},
                     nb_epoch=10,
                     validation_data={'conv_input': val_X, 'skip_input': val_X, 'output': val_Y},
                     batch_size=128,shuffle=True)
    """
    hist = model.fit(X=train_X,y=train_Y,nb_epoch=10,validation_data=(val_X,val_Y),batch_size=128,shuffle=True)
    model.save_weights('metadata/darken_weights_mse_srcnn_C.h5',overwrite=True)
    training = {
        'loss': hist.history['loss'],
        'val_loss': hist.history['val_loss'],
        'optimizer': model.optimizer.get_config(),
    }
    f = open('metadata/darken_training_mse_srcnn_C.json', 'wb')
    meta_json = json.dumps(training, default=lambda o: o.__dict__, indent=4)
    f.write(meta_json)
    f.close()

def inference_darken():
    model = SRCNN_test(img_channels=3,img_rows=512,img_cols=512)
    model.compile(optimizer=SGD(lr=0.01,momentum=0.9,nesterov=True),loss='mse')
    model.load_weights('metadata/darken_weights_mse_srcnn_C.h5')
    IMG_SIZE = 512
    nb_img = 5
    img_folder = '/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/decsai/CM512/inf'
    for img in os.listdir(img_folder):
        print img
        if nb_img <= 0:
            break
        nb_img -= 1
        img_path = img_folder + '/' + img
        img_arr = io.imread(img_path)
        img_arr_darked = exposure.adjust_gamma(image=img_arr,gamma=3)

        img_arr_darked = np.expand_dims(img_arr_darked,0)
        img_arr_darked = img_arr_darked.transpose(0,3,1,2)
        img_recostructed = model.predict(img_arr_darked/255.)

        img_arr_ = np.expand_dims(img_arr,0)
        img_arr_ = img_arr_.transpose(0,3,1,2)
        pred_ori = model.predict(img_arr_/255.)

        """
        input_re = {'conv_input':img_arr_darked/255.,'skip_input':img_arr_darked/255.}
        input_pr = {'conv_input':np.reshape(img_arr,(1,1,img_arr.shape[0],img_arr.shape[1]))/255.
                    ,'skip_input':np.reshape(img_arr,(1,1,img_arr.shape[0],img_arr.shape[1]))/255.}
        img_recostructed = model.predict(data=input_re)['output']
        pred_ori = model.predict(data=input_pr)['output']
        """
        img_recostructed = np.clip(img_recostructed * 255.,0,255)
        pred_ori = np.clip(pred_ori * 255.,0,255)

        img_arr_darked = np.reshape(img_arr_darked.transpose(0,2,3,1),(IMG_SIZE,IMG_SIZE,3))
        img_recostructed = np.reshape(img_recostructed.transpose(0,2,3,1) ,(IMG_SIZE,IMG_SIZE,3))
        pred_ori = np.reshape(pred_ori.transpose(0,2,3,1),(IMG_SIZE,IMG_SIZE,3))
        img_recostructed = util.img_as_ubyte(np.asarray(img_recostructed,dtype=np.uint8))
        pred_ori = util.img_as_ubyte(np.asarray(pred_ori,dtype=np.uint8))
        """
        print img_arr.dtype
        print img_arr_darked.dtype
        print img_recostructed.dtype
        print pred_ori.dtype
        """
        plt.subplot(1,4,1)
        io.imshow(img_arr)
        plt.subplot(1,4,2)
        io.imshow(pred_ori)
        plt.subplot(1,4,3)
        io.imshow(img_arr_darked)
        plt.subplot(1,4,4)
        io.imshow(img_recostructed)
        io.show()
        plt.subplot(1,2,1)
        plt.hist(np.reshape(img_arr,(np.prod(img_arr.shape,))),bins=256)
        plt.subplot(1,2,2)
        plt.hist(np.reshape(pred_ori,(np.prod(pred_ori.shape,))),bins=256)
        plt.show()
        print 'PSNR-originalVSpred_ori: {0}'.format(measure.compare_psnr(im_true=img_arr,im_test=pred_ori,dynamic_range=255))
        print 'PSNR-darkenVSrecostruction: {0}'.format(measure.compare_psnr(im_true=img_arr_darked,im_test=img_recostructed,dynamic_range=255))
        print 'PSNR-originalVSdarken: {0}'.format(measure.compare_psnr(im_true=img_arr,im_test=img_arr_darked,dynamic_range=255))
        print 'PSNR-originalVSrecostruction: {0}'.format(measure.compare_psnr(im_true=img_arr,im_test=img_recostructed,dynamic_range=255))
        print "-"*40
        """
        print 'SSIM-originalVSpred_ori: {0}'.format(measure.compare_ssim(X=img_arr,Y=pred_ori,dynamic_range=255,win_size=7))
        print 'SSIM-noisedVSrecostruction: {0}'.format(measure.compare_ssim(X=img_arr_darked,Y=img_recostructed,dynamic_range=255,win_size=7))
        print 'SSIM-originalVSnoised: {0}'.format(measure.compare_ssim(X=img_arr,Y=img_arr_darked,dynamic_range=255,win_size=7))
        print 'SSIM-originalVSrecostruction: {0}'.format(measure.compare_ssim(X=img_arr,Y=img_recostructed,dynamic_range=255,win_size=7))
        """
        print "*"*40

def inference_noise():
    model = residual_conv(img_channels=1,img_rows=512,img_cols=512)
    model.compile(optimizer=SGD(lr=0.01,momentum=0.9,nesterov=True),loss={'output': 'mse'})
    model.load_weights('metadata/noise_weights_mse_rc.h5')
    nb_img = 5
    img_folder = '/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/decsai/GM512/inf'
    for img in os.listdir(img_folder):
        print img
        if nb_img <= 0:
            break
        nb_img -= 1
        img_path = img_folder + '/' + img
        img_arr = io.imread(img_path)
        img_arr_noised = 255 * util.random_noise(image=img_arr,var=((25/255.)**2)*np.random.uniform(0,1))
        """
        img_recostructed = K.eval(model(K.variable(np.reshape(img_arr_noised,(1,1,img_arr_noised.shape[0],img_arr_noised.shape[1]))/255.)))
        pred_ori = K.eval(model(K.variable(np.reshape(img_arr,(1,1,img_arr.shape[0],img_arr.shape[1]))/255.)))
        """
        input_re = {'conv_input':np.reshape(img_arr_noised,(1,1,img_arr.shape[0],img_arr.shape[1]))/255.
                    ,'skip_input':np.reshape(img_arr_noised,(1,1,img_arr.shape[0],img_arr.shape[1]))/255.}
        input_pr = {'conv_input':np.reshape(img_arr,(1,1,img_arr.shape[0],img_arr.shape[1]))/255.
                    ,'skip_input':np.reshape(img_arr,(1,1,img_arr.shape[0],img_arr.shape[1]))/255.}
        img_recostructed = model.predict(data=input_re)['output']
        pred_ori = model.predict(data=input_pr)['output']
        img_recostructed = np.reshape(img_recostructed,(img_recostructed.shape[2],img_recostructed.shape[3]))
        pred_ori = np.reshape(pred_ori,(pred_ori.shape[2],pred_ori.shape[3]))
        img_recostructed = np.clip(img_recostructed * 255.,0,255)
        pred_ori = np.clip(pred_ori * 255.,0,255)

        img_recostructed = util.img_as_ubyte(np.asarray(img_recostructed,dtype=np.uint8))
        pred_ori = util.img_as_ubyte(np.asarray(pred_ori,dtype=np.uint8))
        img_arr_noised = util.img_as_ubyte(np.asarray(img_arr_noised,dtype=np.uint8))
        # show images
        plt.subplot(1,4,1)
        io.imshow(img_arr)
        plt.subplot(1,4,2)
        io.imshow(pred_ori)
        plt.subplot(1,4,3)
        io.imshow(img_arr_noised)
        plt.subplot(1,4,4)
        io.imshow(img_recostructed)
        plt.show()

        print 'PSNR-originalVSpred_ori: {0}'.format(measure.compare_psnr(im_true=img_arr,im_test=pred_ori,dynamic_range=255))
        print 'PSNR-noisedVSrecostruction: {0}'.format(measure.compare_psnr(im_true=img_arr_noised,im_test=img_recostructed,dynamic_range=255))
        print 'PSNR-originalVSnoised: {0}'.format(measure.compare_psnr(im_true=img_arr,im_test=img_arr_noised,dynamic_range=255))
        print 'PSNR-originalVSrecostruction: {0}'.format(measure.compare_psnr(im_true=img_arr,im_test=img_recostructed,dynamic_range=255))
        """
        print "-"*40
        print 'MMIS-originalVSpred_ori: {0}'.format(measure.compare_mmis(im_true=img_arr,im_test=pred_ori,dynamic_range=255))
        print 'MMIS-noisedVSrecostruction: {0}'.format(measure.compare_mmis(im_true=img_arr_noised,im_test=img_recostructed,dynamic_range=255))
        print 'MMIS-originalVSnoised: {0}'.format(measure.compare_mmis(im_true=img_arr,im_test=img_arr_noised,dynamic_range=255))
        print 'MMIS-originalVSrecostruction: {0}'.format(measure.compare_mmis(im_true=img_arr,im_test=img_recostructed,dynamic_range=255))
        """

        print "*"*40

def inference_darkenANDnoise():
    model = up_pooling_model(img_channels=1,img_rows=512,img_cols=512)
    model.load_weights('metadata/darkenAndnoise_weights_mse_up.h5')
    nb_img = 5
    img_folder = '/media/dell/cb552bf1-c649-4cca-8aca-3c24afca817b/dell/wxm/Data/decsai/GM512/inf'
    for img in os.listdir(img_folder):
        print img
        if nb_img <= 0:
            break
        nb_img -= 1
        img_path = img_folder + '/' + img
        img_arr = io.imread(img_path)
        img_arr_ = exposure.adjust_gamma(image=img_arr,gamma=3)
        img_arr_noised = util.random_noise(image=img_arr_,var=((25/255.)**2)*np.random.uniform(0,1))
        img_arr_noised = np.reshape(img_arr_noised,(1,1,img_arr_noised.shape[0],img_arr_noised.shape[1]))
        pred = K.eval(model({'conv_input':K.variable(img_arr_noised/255.),
                            'skip_input':K.variable(img_arr_noised/255.)}))
        pred_ori = K.eval(model(K.variable(np.reshape(img_arr,(1,1,img_arr.shape[0],img_arr.shape[1]))/255.)))
        pred = np.asarray(pred)
        pred_ori = np.asarray(pred_ori)
        img_recostructed = pred * 255.
        pred_ori = pred_ori * 255.

        img_arr_noised = np.reshape(img_arr_noised,(img_arr.shape[0],img_arr.shape[1]))
        img_recostructed = np.reshape(img_recostructed,(img_arr.shape[0],img_arr.shape[1]))
        pred_ori = np.reshape(pred_ori,(img_arr.shape[0],img_arr.shape[1]))
        #img_recostructed = np.clip(img_recostructed,0,255)
        # show images
        io.imshow(img_arr)
        io.imshow(pred_ori)
        io.imshow(img_arr_noised)
        io.imshow(img_recostructed)
        print 'PSNR-originalVSpred_ori: {0}'.format(PSNR(I=img_arr,K=pred_ori))
        print 'PSNR-noisedVSrecostruction: {0}'.format(PSNR(I=img_arr_noised,K=img_recostructed))
        print 'PSNR-originalVSnoised: {0}'.format(PSNR(I=img_arr,K=img_arr_noised))
        print 'PSNR-originalVSrecostruction: {0}'.format(PSNR(I=img_arr,K=img_recostructed))
        io.show()
        print "*"*40

if __name__ == '__main__':
    inference_darken()